/*
 * Copyright (C) 2025 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "avcodec_codec_name.h"
#include "avcodec_common.h"
#include "plugin/codec_plugin.h"
#include "plugin/plugin_manager_v2.h"
#include <fstream>
#include <gtest/gtest.h>

using namespace std;
using namespace testing::ext;
using namespace OHOS::Media;
using namespace OHOS::Media::Plugins;

namespace OHOS {
namespace Media {
namespace Plugins {
namespace {
const string CODEC_TRUEHD_DEC_NAME = std::string(MediaAVCodec::AVCodecCodecName::AUDIO_DECODER_TRUEHD_NAME);
constexpr int32_t TRUEHD_SAMPLE_RATE = 48000;
constexpr int32_t TRUEHD_CHANNEL_COUNT = 6;
constexpr int32_t TRUEHD_MAX_CHANNEL_COUNT = 8;
constexpr int64_t TRUEHD_BIT_RATE = 18000000;
constexpr int32_t TRUEHD_MAX_INPUT_SIZE = 24576;
} // namespace

class TRUEHDUnitTest : public testing::Test, public DataCallback {
public:
    static void SetUpTestCase(void);
    static void TearDownTestCase(void);
    void SetUp() override;
    void TearDown() override;

    void OnInputBufferDone(const shared_ptr<Media::AVBuffer> &inputBuffer) override
    {
        (void)inputBuffer;
    }
    void OnOutputBufferDone(const shared_ptr<Media::AVBuffer> &outputBuffer) override
    {
        (void)outputBuffer;
    }
    void OnEvent(const shared_ptr<Plugins::PluginEvent> event) override
    {
        (void)event;
    }

protected:
    shared_ptr<CodecPlugin> CreatePlugin();
    shared_ptr<Media::Meta> CreateMeta();
    shared_ptr<AVBuffer> CreateAVBuffer(size_t capacity, size_t size);
    shared_ptr<CodecPlugin> plugin_ = nullptr;
    shared_ptr<Media::Meta> meta_ = nullptr;
};

void TRUEHDUnitTest::SetUpTestCase(void) {}

void TRUEHDUnitTest::TearDownTestCase(void) {}

void TRUEHDUnitTest::SetUp(void)
{
    plugin_ = CreatePlugin();
    ASSERT_NE(plugin_, nullptr);
    meta_ = CreateMeta();
    ASSERT_NE(meta_, nullptr);
}

void TRUEHDUnitTest::TearDown(void)
{
    if (plugin_) {
        plugin_->Release();
    }
}

shared_ptr<CodecPlugin> TRUEHDUnitTest::CreatePlugin()
{
    auto plugin = PluginManagerV2::Instance().CreatePluginByName(CODEC_TRUEHD_DEC_NAME);
    if (!plugin) {
        return nullptr;
    }
    return reinterpret_pointer_cast<CodecPlugin>(plugin);
}

shared_ptr<Media::Meta> TRUEHDUnitTest::CreateMeta()
{
    auto meta = make_shared<Media::Meta>();
    meta->Set<Tag::AUDIO_CHANNEL_COUNT>(TRUEHD_CHANNEL_COUNT);
    meta->Set<Tag::AUDIO_SAMPLE_RATE>(TRUEHD_SAMPLE_RATE);
    meta->Set<Tag::AUDIO_CHANNEL_LAYOUT>(AudioChannelLayout::CH_5POINT1);
    meta->Set<Tag::MEDIA_BITRATE>(TRUEHD_BIT_RATE);
    meta->Set<Tag::AUDIO_SAMPLE_FORMAT>(AudioSampleFormat::SAMPLE_S32LE);
    return meta;
}

shared_ptr<AVBuffer> TRUEHDUnitTest::CreateAVBuffer(size_t capacity, size_t size)
{
    auto avAllocator = AVAllocatorFactory::CreateSharedAllocator(MemoryFlag::MEMORY_READ_WRITE);
    auto avBuffer = AVBuffer::CreateAVBuffer(avAllocator, capacity);
    avBuffer->memory_->SetSize(size);
    return avBuffer;
}

HWTEST_F(TRUEHDUnitTest, SetParameter_001, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
}

HWTEST_F(TRUEHDUnitTest, SetParameter_002, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    meta_->Set<Tag::AUDIO_CHANNEL_COUNT>(0);
    ASSERT_NE(plugin_->SetParameter(meta_), Status::OK);
    meta_->Set<Tag::AUDIO_CHANNEL_COUNT>(-1);
    ASSERT_NE(plugin_->SetParameter(meta_), Status::OK);
    meta_->Set<Tag::AUDIO_CHANNEL_COUNT>(TRUEHD_CHANNEL_COUNT);
    meta_->Set<Tag::AUDIO_SAMPLE_RATE>(0);
    ASSERT_NE(plugin_->SetParameter(meta_), Status::OK);
    meta_->Set<Tag::AUDIO_SAMPLE_RATE>(-1);
    ASSERT_NE(plugin_->SetParameter(meta_), Status::OK);
    meta_->Set<Tag::AUDIO_SAMPLE_RATE>(TRUEHD_SAMPLE_RATE);
}

HWTEST_F(TRUEHDUnitTest, SetParameter_003, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    meta_->Remove(Tag::AUDIO_CHANNEL_COUNT);
    ASSERT_NE(plugin_->SetParameter(meta_), Status::OK);
    meta_->Set<Tag::AUDIO_CHANNEL_COUNT>(TRUEHD_CHANNEL_COUNT);
    meta_->Remove(Tag::AUDIO_SAMPLE_RATE);
    ASSERT_NE(plugin_->SetParameter(meta_), Status::OK);
    meta_->Set<Tag::AUDIO_SAMPLE_RATE>(TRUEHD_SAMPLE_RATE);
}

HWTEST_F(TRUEHDUnitTest, Lifecycle_001, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->Start(), Status::OK);
    ASSERT_EQ(plugin_->Flush(), Status::OK);
    ASSERT_EQ(plugin_->Stop(), Status::OK);
    ASSERT_EQ(plugin_->Reset(), Status::OK);
    ASSERT_EQ(plugin_->Release(), Status::OK);
}

HWTEST_F(TRUEHDUnitTest, Eos_001, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetDataCallback(this), Status::OK);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->Start(), Status::OK);
    auto inputBuffer = CreateAVBuffer(1024, 0);
    inputBuffer->flag_ = MediaAVCodec::AVCODEC_BUFFER_FLAG_EOS;
    ASSERT_EQ(plugin_->QueueInputBuffer(inputBuffer), Status::OK);
    auto outputBuffer = CreateAVBuffer(TRUEHD_MAX_INPUT_SIZE * 4, 0);
    ASSERT_EQ(plugin_->QueueOutputBuffer(outputBuffer), Status::END_OF_STREAM);
    ASSERT_EQ(plugin_->Stop(), Status::OK);
}

HWTEST_F(TRUEHDUnitTest, GetParameter_001, TestSize.Level1)
{
    meta_->Set<Tag::AUDIO_MAX_INPUT_SIZE>(TRUEHD_MAX_INPUT_SIZE);
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    shared_ptr<Media::Meta> outMeta = make_shared<Media::Meta>();
    ASSERT_EQ(plugin_->GetParameter(outMeta), Status::OK);
    int32_t channelCount = 0;
    int32_t sampleRate = 0;
    int32_t maxInputSize = 0;
    ASSERT_TRUE(outMeta->Get<Tag::AUDIO_CHANNEL_COUNT>(channelCount));
    ASSERT_EQ(channelCount, TRUEHD_CHANNEL_COUNT);
    ASSERT_TRUE(outMeta->Get<Tag::AUDIO_SAMPLE_RATE>(sampleRate));
    ASSERT_EQ(sampleRate, TRUEHD_SAMPLE_RATE);
    ASSERT_TRUE(outMeta->Get<Tag::AUDIO_MAX_INPUT_SIZE>(maxInputSize));
    ASSERT_GT(maxInputSize, 0);
}

HWTEST_F(TRUEHDUnitTest, GetParameter_002, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetDataCallback(this), Status::OK);
    int32_t defaultInputSize = 7680 * TRUEHD_MAX_CHANNEL_COUNT * 4;
    shared_ptr<Media::Meta> outMeta = make_shared<Media::Meta>();
    int32_t maxInputSize = 0;
    int32_t smallInputSize = 4096;
    meta_->Set<Tag::AUDIO_MAX_INPUT_SIZE>(smallInputSize);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->GetParameter(outMeta), Status::OK);
    ASSERT_TRUE(outMeta->Get<Tag::AUDIO_MAX_INPUT_SIZE>(maxInputSize));
    ASSERT_EQ(maxInputSize, smallInputSize);
    int32_t largeInputSize = 10000000;
    meta_->Set<Tag::AUDIO_MAX_INPUT_SIZE>(largeInputSize);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->GetParameter(outMeta), Status::OK);
    ASSERT_TRUE(outMeta->Get<Tag::AUDIO_MAX_INPUT_SIZE>(maxInputSize));
    ASSERT_EQ(maxInputSize, defaultInputSize);
    meta_->Set<Tag::AUDIO_MAX_INPUT_SIZE>(-1);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->GetParameter(outMeta), Status::OK);
    ASSERT_TRUE(outMeta->Get<Tag::AUDIO_MAX_INPUT_SIZE>(maxInputSize));
    ASSERT_EQ(maxInputSize, defaultInputSize);
}

HWTEST_F(TRUEHDUnitTest, Decode_With_Invalid_File_001, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetDataCallback(this), Status::OK);
    meta_->Set<Tag::AUDIO_MAX_INPUT_SIZE>(TRUEHD_MAX_INPUT_SIZE);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->Start(), Status::OK);
    auto inputBuffer = CreateAVBuffer(TRUEHD_MAX_INPUT_SIZE, 1024);
    auto outputBuffer = CreateAVBuffer(TRUEHD_MAX_INPUT_SIZE * 4, 0);
    // No valid data, so it errors
    ASSERT_EQ(plugin_->QueueInputBuffer(inputBuffer), Status::ERROR_INVALID_DATA);
    Status status = plugin_->QueueOutputBuffer(outputBuffer);
    ASSERT_TRUE(status == Status::OK || status == Status::ERROR_NOT_ENOUGH_DATA || status == Status::ERROR_UNKNOWN);
    ASSERT_EQ(plugin_->Stop(), Status::OK);
}

HWTEST_F(TRUEHDUnitTest, Decode_With_Valid_File_001, TestSize.Level1)
{
    ASSERT_EQ(plugin_->Init(), Status::OK);
    ASSERT_EQ(plugin_->SetDataCallback(this), Status::OK);
    meta_->Set<Tag::AUDIO_MAX_INPUT_SIZE>(TRUEHD_MAX_INPUT_SIZE);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    ASSERT_EQ(plugin_->Start(), Status::OK);
    const std::string truehdFilePath = "/data/test/media/truehd_44100_stereo.dat";
    std::ifstream inFile(truehdFilePath, std::ios::binary);
    ASSERT_TRUE(inFile.is_open());
    uint32_t channelCount = 0, sampleRate = 0, extraDataLen = 0;
    inFile.read(reinterpret_cast<char *>(&channelCount), sizeof(channelCount));
    inFile.read(reinterpret_cast<char *>(&sampleRate), sizeof(sampleRate));
    inFile.read(reinterpret_cast<char *>(&extraDataLen), sizeof(extraDataLen));
    std::vector<uint8_t> extraData(extraDataLen);
    if (extraDataLen > 0) {
        inFile.read(reinterpret_cast<char *>(extraData.data()), extraDataLen);
    }
    meta_->Set<Tag::AUDIO_CHANNEL_COUNT>(channelCount);
    meta_->Set<Tag::AUDIO_SAMPLE_RATE>(sampleRate);
    meta_->Set<Tag::AUDIO_CHANNEL_LAYOUT>(AudioChannelLayout::STEREO);
    ASSERT_EQ(plugin_->SetParameter(meta_), Status::OK);
    while (inFile.good()) {
        uint64_t frameSize = 0, pts = 0;
        inFile.read(reinterpret_cast<char *>(&frameSize), sizeof(frameSize));
        inFile.read(reinterpret_cast<char *>(&pts), sizeof(pts));
        if (inFile.gcount() != sizeof(pts)) {
            break;
        }
        if (frameSize == 0 || frameSize > TRUEHD_MAX_INPUT_SIZE) {
            inFile.seekg(frameSize, std::ios::cur);
            continue;
        }
        std::vector<uint8_t> frameData(frameSize);
        inFile.read(reinterpret_cast<char *>(frameData.data()), frameSize);
        if (inFile.gcount() != static_cast<std::streamsize>(frameSize)) {
            break;
        }
        auto inputBuffer = CreateAVBuffer(TRUEHD_MAX_INPUT_SIZE, frameSize);
        std::copy(frameData.data(), frameData.data() + frameSize,
            static_cast<uint8_t*>(inputBuffer->memory_->GetAddr()));
        inputBuffer->memory_->SetSize(frameSize);
        inputBuffer->flag_ = 0;
        inputBuffer->pts_ = pts;
        Status ret = plugin_->QueueInputBuffer(inputBuffer);
        ASSERT_TRUE(ret == Status::OK || ret == Status::ERROR_NOT_ENOUGH_DATA);
    }
    inFile.close();
    auto eosBuffer = CreateAVBuffer(1024, 0);
    eosBuffer->flag_ = MediaAVCodec::AVCODEC_BUFFER_FLAG_EOS;
    ASSERT_EQ(plugin_->QueueInputBuffer(eosBuffer), Status::OK);
    auto outputBuffer = CreateAVBuffer(TRUEHD_MAX_INPUT_SIZE * 4, 0);
    Status status = plugin_->QueueOutputBuffer(outputBuffer);
    ASSERT_TRUE(status == Status::OK || status == Status::END_OF_STREAM || status == Status::ERROR_NOT_ENOUGH_DATA ||
                status == Status::ERROR_AGAIN);
    ASSERT_EQ(plugin_->Stop(), Status::OK);
}

} // namespace Plugins
} // namespace Media
} // namespace OHOS
